using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using UnityEngine;

namespace ZerAx.Tweaks.Collections
{
    /// <summary>
    /// A reference-counted and array-like container for complicated but repeated situations.
    /// 
    /// Inspired by Flyweight pattern: https://gameprogrammingpatterns.com/flyweight.html
    /// </summary>
    /// <typeparam name="T">The value you want contain.</typeparam>
    public class Palette3D<T> where T : class
    {
        [Flags]
        public enum PaletteStatus : byte
        {
            Fine = 0,
            NeedClean = 1,
            ValuesChanged = 2,
        }

        public readonly int lenX, lenY, lenZ;
        public int?[] indices;
        public List<Rc<T>> values;

        public IEnumerable<T> Values => values.Select(x => x.Deref());
        public IEnumerable<Rc<T>> Entries => values.Where(x => !x.CanDestroy);
        public PaletteStatus Status { get; private set; }
        public bool IsDirty => Status.HasFlag(PaletteStatus.ValuesChanged);

        public T this[int x, int y, int z]
        {
            get => Get(x, y, z);
            set => Set(x, y, z, value);
        }

        public Palette3D(int lenX, int lenY, int lenZ)
        {
            indices = new int?[lenX * lenY * lenZ];
            this.lenX = lenX;
            this.lenY = lenY;
            this.lenZ = lenZ;
            values = new List<Rc<T>>();
            Status = PaletteStatus.Fine;
        }

        public int InnerIndex(int x, int y, int z) => z * lenX * lenY + y * lenX + x;

        public int? GetIndex(int x, int y, int z) => indices[InnerIndex(x, y, z)];

        public T Get(int x, int y, int z)
        {
            if (GetIndex(x, y, z) is int idx)
            {
                return values[idx].Deref();
            }
            return null;
        }

        public void Set(int x, int y, int z, T value)
        {
            if (value == null)
            {
                Remove(x, y, z);
                return;
            }

            Drop(x, y, z);
            int idx = values.FindIndex(x => x.Deref() == value);
            if (idx < 0)
            {
                indices[InnerIndex(x, y, z)] = values.Count;
                values.Add(new Rc<T>(value));
                Status |= PaletteStatus.ValuesChanged;
            }
            else
            {
                values[idx].MakeRef();
                indices[InnerIndex(x, y, z)] = idx;
            }

            if (Status.HasFlag(PaletteStatus.NeedClean)) Clean();
        }

        public void Remove(int x, int y, int z)
        {
            Drop(x, y, z);

            if (Status.HasFlag(PaletteStatus.NeedClean)) Clean();
        }

        private void Drop(int x, int y, int z)
        {
            Drop(InnerIndex(x, y, z));
        }

        private void Drop(int innerIndex)
        {
            if (indices[innerIndex] is int idx)
            {
                indices[innerIndex] = null;
                values[idx].DropRef();
                if (values[idx].CanDestroy) Status |= PaletteStatus.NeedClean;
            }
        }

        public void BatchedRemove(IEnumerable<(int, int, int)> batch)
        {
            foreach (var (x, y, z) in batch)
            {
                Drop(x, y, z);
            }
            if (Status.HasFlag(PaletteStatus.NeedClean)) Clean();
        }

        public void BatchedSet(IEnumerable<(int, int, int)> batch, T value)
        {
            if (value == null)
            {
                BatchedRemove(batch); 
                return;
            }

            int idx = values.FindIndex(x => x.Deref() == value);
            if (idx < 0)
            {
                Rc<T> entry = new(value, 0);
                foreach (var (x, y, z) in batch)
                {
                    Drop(x, y, z);
                    entry.MakeRef();
                    indices[InnerIndex(x, y, z)] = values.Count;
                }
                values.Add(entry);
                Status |= PaletteStatus.ValuesChanged;
            }
            else
            {
                var entry = values[idx];
                foreach (var (x, y, z) in batch)
                {
                    Drop(x, y, z);
                    entry.MakeRef();
                    indices[InnerIndex(x, y, z)] = idx;
                }
                values[idx] = entry;
            }

            if (Status.HasFlag(PaletteStatus.NeedClean)) Clean();
        }

        public void Clean()
        {
            var mapping = new int?[values.Count];
            var cleanedValues = Entries.ToArray();

            for (int i = 0; i < mapping.Length; i++)
            {
                if (values[i].CanDestroy)
                {
                    mapping[i] = null;
                }
                else
                {
                    for (int j = 0; j < cleanedValues.Length; j++)
                    {
                        if (cleanedValues[j].Deref() == values[i].Deref())
                        {
                            mapping[i] = j;
                            break;
                        }
                    }
                }
            }

            for (int i = 0; i < indices.Length; i++)
            {
                if (indices[i] is int index) indices[i] = mapping[index];
            }
            values.Clear();
            values.AddRange(cleanedValues);
            
            Status &= ~PaletteStatus.NeedClean;
            Status |= PaletteStatus.ValuesChanged;
        }

        public void Clear()
        {
            for (int i = 0; i < indices.Length; i++)
            {
                indices[i] = null;
            }
            values.Clear();
            Status |= PaletteStatus.ValuesChanged;
        }

        public void UnmarkDirty()
        {
            Status = PaletteStatus.Fine;
        }

        public class Rc<Ty>
        {
            public Ty val;
            public int refCount;

            public Rc(Ty val, int refCount = 1)
            {
                this.val = val;
                this.refCount = refCount;
            }

            public Ty Deref() => val;
            public void MakeRef() => refCount += 1;
            public void DropRef() => refCount -= 1;
            public bool CanDestroy => refCount <= 0;

            public override string ToString() => $"<{val},{refCount}>";
        }
    }
}
